{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%html\n",
    "<style>\n",
    ".cell-output-ipywidget-background {\n",
    "    background-color: transparent !important;\n",
    "}\n",
    ":root {\n",
    "    --jp-widgets-color: var(--vscode-editor-foreground);\n",
    "    --jp-widgets-font-size: var(--vscode-editor-font-size);\n",
    "}  \n",
    "</style>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import nest_asyncio\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# PREPARE ENVIRONMENT\n",
    "os.environ[\"MODAL_IMAGE_BUILDER_VERSION\"] = \"2024.10\"\n",
    "os.environ[\"SWE_AGENT_CONFIG_DIR\"] = \".\"\n",
    "os.environ[\"SWE_AGENT_TOOLS_DIR\"] = \"tools\"\n",
    "os.environ[\"SWE_AGENT_TRAJECTORY_DIR\"] = \"trajectories\"\n",
    "os.environ[\"WEAVE_PRINT_CALL_LINK\"] = \"False\"\n",
    "os.environ[\"WEAVE_LOG_LEVEL\"] = \"CRITICAL\"\n",
    "\n",
    "os.makedirs(\"replays\", exist_ok=True)\n",
    "os.makedirs(\"trajectories\", exist_ok=True)\n",
    "\n",
    "load_dotenv()\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from rollout import ModelConfig\n",
    "\n",
    "import art\n",
    "from art.local import LocalBackend\n",
    "\n",
    "# INITIALIZE MODEL\n",
    "backend = LocalBackend()\n",
    "model = art.TrainableModel(\n",
    "    name=\"031\",\n",
    "    project=\"sweagent\",\n",
    "    config=ModelConfig(\n",
    "        max_input_tokens=40_960,\n",
    "        system_prompt_suffix=\"\\n/no_think\",\n",
    "        xml_function_calling=True,\n",
    "    ),\n",
    "    base_model=\"Qwen/Qwen3-32B\",\n",
    "    _internal_config=art.dev.InternalModelConfig(\n",
    "        engine_args=art.dev.EngineArgs(\n",
    "            tensor_parallel_size=torch.cuda.device_count(), gpu_memory_utilization=0.85\n",
    "        ),\n",
    "        torchtune_args=art.dev.TorchtuneArgs(\n",
    "            model=\"qwen3_32b\", model_type=\"QWEN3\", async_weight_syncing=True\n",
    "        ),\n",
    "    ),\n",
    ")\n",
    "await model.register(backend)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sandboxes import terminate_sandboxes\n",
    "\n",
    "await terminate_sandboxes()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import weave\n",
    "\n",
    "from instances import as_instances_iter, get_filtered_swe_smith_instances_df\n",
    "from rollout import rollout\n",
    "\n",
    "# weave.init(project_name=model.project)\n",
    "# rollout = weave.op(rollout)\n",
    "\n",
    "# TRAIN MODEL\n",
    "instances = list(\n",
    "    get_filtered_swe_smith_instances_df()\n",
    "    .sample(fraction=1.0, shuffle=True, seed=42)\n",
    "    .pipe(as_instances_iter)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "GROUPS_PER_BATCH = 16\n",
    "ROLLOUTS_PER_GROUP = 16\n",
    "# for i in range(await model.get_step(), 1_000):\n",
    "for i in [0] * 1_000:\n",
    "    trajectory_groups = await art.gather_trajectory_groups(\n",
    "        (\n",
    "            art.TrajectoryGroup(\n",
    "                rollout(model, instance, reward_power=1.33, timeout=60 * 15)\n",
    "                for _ in range(ROLLOUTS_PER_GROUP)\n",
    "            )\n",
    "            for instance in instances[\n",
    "                i * GROUPS_PER_BATCH % len(instances) : (i + 1)\n",
    "                * GROUPS_PER_BATCH\n",
    "                % len(instances)\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    await model.train(\n",
    "        trajectory_groups,\n",
    "        # config=art.TrainConfig(learning_rate=7e-6),\n",
    "        _config=art.dev.TrainConfig(allow_training_without_logprobs=True),\n",
    "        verbose=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "from typing import Any, AsyncIterable, Coroutine, Iterable, TypeVar\n",
    "\n",
    "T = TypeVar(\"T\")\n",
    "\n",
    "\n",
    "async def test() -> int: ...\n",
    "\n",
    "\n",
    "d = test()\n",
    "\n",
    "\n",
    "async def buffer(\n",
    "    iterable: Iterable[Coroutine[Any, Any, T]], *, max_concurrent: int\n",
    ") -> AsyncIterable[list[T]]:\n",
    "    buffer: list[asyncio.Task[T]] = []\n",
    "    iterator = iter(iterable)\n",
    "    iterator_exhausted = False\n",
    "\n",
    "    while buffer or not iterator_exhausted:\n",
    "        # Fill buffer up to max_concurrent\n",
    "        while (\n",
    "            len([task for task in buffer if not task.done()]) < max_concurrent\n",
    "            and not iterator_exhausted\n",
    "        ):\n",
    "            try:\n",
    "                buffer.append(asyncio.create_task(next(iterator)))\n",
    "            except StopIteration:\n",
    "                iterator_exhausted = True\n",
    "                break\n",
    "        if not buffer:\n",
    "            break\n",
    "        await asyncio.wait(buffer, return_when=asyncio.FIRST_COMPLETED)\n",
    "        for task in buffer:\n",
    "            if task.done() and (exception := task.exception()):\n",
    "                raise exception\n",
    "        results = []\n",
    "        for task in list(buffer):\n",
    "            if task.done():\n",
    "                results.append(task.result())\n",
    "                buffer.remove(task)\n",
    "            else:\n",
    "                break\n",
    "        if results:\n",
    "            yield results\n",
    "\n",
    "\n",
    "GROUPS_PER_BATCH = 4\n",
    "ROLLOUTS_PER_GROUP = 4\n",
    "\n",
    "# Calculate the number of digits needed for batch numbering\n",
    "total_batches = (len(instances) + GROUPS_PER_BATCH - 1) // GROUPS_PER_BATCH\n",
    "num_digits = len(str(total_batches - 1))\n",
    "\n",
    "async for trajectory_groups in buffer(\n",
    "    (\n",
    "        art.gather_trajectory_groups(\n",
    "            (\n",
    "                art.TrajectoryGroup(\n",
    "                    rollout(model, instance) for _ in range(ROLLOUTS_PER_GROUP)\n",
    "                )\n",
    "                for instance in instances[start : start + GROUPS_PER_BATCH]\n",
    "            ),\n",
    "            pbar_desc=f\"gather/{start // GROUPS_PER_BATCH:0{num_digits}d}\",\n",
    "        )\n",
    "        for start in range(0, len(instances), GROUPS_PER_BATCH)\n",
    "    ),\n",
    "    max_concurrent=3,\n",
    "):\n",
    "    await model.train(\n",
    "        [g for gs in trajectory_groups for g in gs],\n",
    "        _config=art.dev.TrainConfig(allow_training_without_logprobs=True),\n",
    "    )\n",
    "\n",
    "\n",
    "# async for trajectory_groups in art.trajectory_group_batches(\n",
    "#     (\n",
    "#         art.TrajectoryGroup(rollout(model, instance) for _ in range(4))\n",
    "#         for instance in instances\n",
    "#     ),\n",
    "#     batch_size=4,\n",
    "#     max_concurrent_batches=3,\n",
    "#     skip_batches=await model.get_step(),\n",
    "# ):\n",
    "#     await model.train(\n",
    "#         trajectory_groups,\n",
    "#         _config=art.dev.TrainConfig(allow_training_without_logprobs=True),\n",
    "#     )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
